# -*- coding: utf-8 -*-
"""
@Env 
@Time 2024/9/10 上午10:15
@Author yzpang
@Function: 子进程模型加载
"""
import os
import time
from multiprocessing import Process


SYSTEM = os.name
# 模型文件绝对路径
if SYSTEM == 'nt':
    MODEL_ROOT_PATH = r"E:\AI\models"
    DATASET_ROOT_PATH = r"E:\AI\datasets"
else:
    MODEL_ROOT_PATH = r"/home/models"
    DATASET_ROOT_PATH = r"/home/datasets"


def load_pre_model(model_path):
    from transformers import AutoModelForSequenceClassification, AutoTokenizer
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForSequenceClassification.from_pretrained(model_path)
        while True:
            time.sleep(1)
    except Exception as e:
        print(f"Exception: {e}")


def load_sub_process(model_path, **kwargs):
    process = Process(target=load_pre_model, args=(model_path,))
    process.start()


def load_model():
    model_path = os.path.join(MODEL_ROOT_PATH, "BERT/bert-base-chinese")
    load_sub_process(model_path)


if __name__ == '__main__':
    load_model()